package com.dyrnq.seckill.interceptor;

import io.github.bucket4j.*;
import io.github.bucket4j.distributed.ExpirationAfterWriteStrategy;
import io.github.bucket4j.redis.lettuce.cas.LettuceBasedProxyManager;
import io.lettuce.core.RedisClient;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.codec.ByteArrayCodec;
import io.lettuce.core.codec.RedisCodec;
import io.lettuce.core.codec.StringCodec;
import lombok.AllArgsConstructor;
import lombok.RequiredArgsConstructor;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.time.Duration;

import static java.time.Duration.ofSeconds;


@Component
@RequiredArgsConstructor
public class RateLimitInterceptor implements HandlerInterceptor {


    final RedisClient redisClient;


    private Bucket resolveBucket(String key) {


        StatefulRedisConnection<String, byte[]> redisConnection = redisClient.connect(RedisCodec.of(StringCodec.UTF8, ByteArrayCodec.INSTANCE));
        LettuceBasedProxyManager<String> proxyManager = LettuceBasedProxyManager.builderFor(redisConnection)
                .withExpirationStrategy(ExpirationAfterWriteStrategy.basedOnTimeForRefillingBucketUpToMax(ofSeconds(10)))
                .build();

        Refill refill = Refill.intervally(100, Duration.ofMinutes(1));
        Bandwidth limit = Bandwidth.classic(100, refill);

        BucketConfiguration configuration = BucketConfiguration.builder()
                .addLimit(limit)
                .build();
        Bucket bucket = proxyManager.builder().build(key, configuration);

        return bucket;
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
            throws Exception {
//        String apiKey = request.getHeader("X-api-key");
//        if (apiKey == null || apiKey.isEmpty()) {
//            response.sendError(HttpStatus.BAD_REQUEST.value(), "Missing Header: X-api-key");
//            return false;
//        }
        String apiKey = request.getRequestURI().toLowerCase();

        String user_id = null;
        try {
            user_id = request.getParameter("user_id");

        } catch (Exception e) {

        }

        if (user_id != null) {
            apiKey = apiKey + ":" + user_id;
        }


        Bucket tokenBucket = resolveBucket(apiKey);
        ConsumptionProbe probe = tokenBucket.tryConsumeAndReturnRemaining(1);
        if (probe.isConsumed()) {
            response.addHeader("X-Rate-Limit-Remaining", String.valueOf(probe.getRemainingTokens()));
            return true;
        } else {
            long waitForRefill = probe.getNanosToWaitForRefill() / 1_000_000_000;
            response.addHeader("X-Rate-Limit-Retry-After-Seconds", String.valueOf(waitForRefill));
            response.sendError(HttpStatus.TOO_MANY_REQUESTS.value(),
                    "You have exhausted your API Request Quota");
            return false;
        }
    }
}
